module nh_case1_mod
  ! Steady flow
  use const_mod

  implicit none

  private

  public nh_case1

  real(8), parameter :: Rd      = 287.04
  real(8), parameter :: Rv      = 461.5
  real(8), parameter :: cp      = 1004.5
  real(8), parameter :: cv      = 717.5
  real(8), parameter :: omega   = 7.292d-5
  real(8), parameter :: p0      = 100000
  real(8), parameter :: ps      = p0
  real(8), parameter :: p_top   = 0
  real(8), parameter :: t0      = 288
  real(8), parameter :: u0      = 35
  real(8), parameter :: eta_s   = 1
  real(8), parameter :: eta_t   = 0.2
  real(8), parameter :: eta0    = 0.252
  real(8), parameter :: gamma   = 0.005
  real(8), parameter :: delta_t = 4.8e5
  real(8), parameter :: radius  = 6371229
  real(8), parameter :: gravity = 9.80616d0

contains

  subroutine nh_case1(lon, lat, z, rho, u, v, w, pt, mr, zs, rc, g, omg, radi, zt)
    real(8), intent(in )           :: lon
    real(8), intent(in )           :: lat
    real(8), intent(in ), optional :: z
    real(8), intent(out), optional :: rho
    real(8), intent(out), optional :: u
    real(8), intent(out), optional :: v
    real(8), intent(out), optional :: w
    real(8), intent(out), optional :: pt
    real(8), intent(out), optional :: mr
    real(8), intent(out), optional :: zs
    real(8), intent(out), optional :: rc ! Rayleigh damping coefficient
    real(8), intent(out), optional :: g
    real(8), intent(out), optional :: omg
    real(8), intent(out), optional :: radi
    real(8), intent(out), optional :: zt
    
    real(8) :: t, p, eta, eta_v, f, dfdeta, phi, phi_bar, t_bar
    
    integer iter
    real(8) :: eta_prev
    real(8) :: residual
    
    if( ( present(rho) .or. present(u) .or. present(v) .or. present(w) .or. present(pt) .or. present(mr) .or. present(rc) ) &
    .and. .not. present(z) )then
      stop 'Need z in nh_case1 while acquiring rho, u, v, w, pt, mr, rc'
    endif
    
    if( present(z) )then
      eta = eta_t
      iter = 0
      residual = 10
      
      do while( residual > 1.e-14 )
        iter   = iter + 1
        t      = calc_t( eta, lat )
        phi    = calc_phi( eta, lat )
        f      = phi - gravity * z
        dfdeta = - Rd / eta * t
        
        eta_prev = eta
        eta      = eta - f / dfdeta
        residual = abs( eta - eta_prev )
      enddo
      p     = p0 * eta
      eta_v = ( eta - eta0 ) * pi / 2.
    endif
    
    if (present(radi  )) radi  = radius
    if (present(rho   )) rho   = p / ( Rd * t )
    if (present(u     )) u     = u0 * cos(eta_v)**1.5 * sin(2*lat)**2
    if (present(v     )) v     = 0
    if (present(w     )) w     = 0
    if (present(pt    )) pt    = t * ( p0 / p )**( Rd / Cp )
    if (present(mr    )) mr    = 0
    if (present(zs    )) zs    = calc_topo( lat )
    if (present(rc    )) rc    = 0
    if (present(g     )) g     = gravity
    if (present(omg   )) omg   = omega
    if (present(zt    )) zt    = calc_phi( eta0, lat ) / gravity

  end subroutine nh_case1
  
  function calc_t( eta, lat ) result (t)
    real(8) :: t
    real(8) :: eta
    real(8) :: lat
    
    real(8) :: t_bar, eta_v
    
    eta_v = ( eta - eta0 ) * pi / 2.
    t_bar = calc_t_bar( eta )
    t  = t_bar + 0.75 * eta * pi * u0 / Rd * sin(eta_v) * sqrt(cos(eta_v)) &
               * ( ( -2 * sin(lat)**6 * ( cos(lat)**2 + 1./3. ) + 10./63. ) * 2 * u0 * cos(eta_v)**1.5 &
               + ( 8./5. * cos(lat)**3 * ( sin(lat)**2 + 2./3. ) - pi/4. ) * radius * omega )
    
  end function calc_t
  
  function calc_phi( eta, lat ) result (phi)
    real(8) :: phi
    real(8) :: eta
    real(8) :: lat
    
    real(8) :: phi_bar, eta_v
    
    eta_v = ( eta - eta0 ) * pi / 2.
    phi_bar = calc_phi_bar( eta )
    phi = phi_bar + u0 * cos(eta_v)**1.5                                                    &
        * ( ( -2 * sin(lat)**6 * ( cos(lat)**2 + 1./3. ) + 10./63. ) * u0 * cos(eta_v)**1.5 &
        + ( 8./5. * cos(lat)**3 * ( sin(lat)**2 + 2./3. ) - pi/4. ) * radius * omega )
    
  end function calc_phi
  
  function calc_t_bar( eta ) result (t_bar)
    real(8) :: t_bar
    real(8) :: eta
    
    real(8) :: g
    
    g = gravity
    
    t_bar  = merge( t0 * eta**(Rd*gamma/g), &
                    t0 * eta**(Rd*gamma/g) + delta_t * ( eta_t - eta )**5, eta >= eta_t )
    
  end function calc_t_bar
  
  function calc_phi_bar( eta ) result (phi_bar)
    real(8) :: phi_bar
    real(8) :: eta
    
    real(8) :: g
    
    g = gravity
    
    phi_bar  =  merge( t0 * g / gamma * ( 1. - eta**(Rd*gamma/g) ), &
                       t0 * g / gamma * ( 1. - eta**(Rd*gamma/g) )  &
                       - Rd * delta_t *                             &
                       ( ( log(eta/eta_t) + 137./60. ) * eta_t**5   &
                       - 5 * eta_t**4 * eta                         &
                       + 5 * eta_t**3 * eta**2                      &
                       - 10./3. * eta_t**2 * eta**3                 &
                       + 5./4. * eta_t * eta**4                     &
                       - 0.2 * eta**5 ),                            &
                       eta >= eta_t)
    
  end function calc_phi_bar
  
  function calc_topo( lat ) result( zs )
    real(8) :: lat
    real(8) :: zs
    
    real(8) :: eta, eta_v
    
    eta   = eta_s
    eta_v = ( eta - eta0 ) * pi / 2.
    
    zs = u0 * cos(eta_v)**1.5                                                              &
       * ( ( -2 * sin(lat)**6 * ( cos(lat)**2 + 1./3. ) + 10./63. ) * u0 * cos(eta_v)**1.5 &
       + ( 8./5. * cos(lat)**3 * ( sin(lat)**2 + 2./3. ) - pi/4. ) * radius * omega )
    
    zs = zs / gravity
    
  end function calc_topo
  
end module nh_case1_mod
